# This program computes and plots the TM mode fields in
# 2D rectagnular space with PEC boundaries using
# Yee's FDTD Algorithm.
#
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.animation as animation
# from IPython.display import HTML

## Define Parameters
c = 2.99792458e8  # Speed of light
eps0 = 8.854187817*10**(-12)  # Free space permittivity
xmu0 = 4 * math.pi *10**(-7) # Free space permeability
#
## Define Space Region
a = 10
b = 10  # Space Dimensions
# Number of grid points in x and y direction
nx = 100
ny = 100
# Convert to integer in‑case of odd number
nx2 = math.ceil(nx / 2)
ny2 = math.ceil(ny / 2)
nt = 200  # Number of time‑steps
nskip = 5  # Time‑steps to skip while recording
nsnap = math.ceil(nt / nskip)

#
## Create a Gaussian Excitation Source
xndec = 10.0  # Mean
xn0 = 4 * xndec  # Variance

#
## Initialize

Ez = np.zeros((nx, ny))  # Z E‑field initialize to zero
Hx = np.zeros((nx, ny))  # X H‑field initialize to zero
Hy = np.zeros((nx, ny))  # Y H‑field initialize to zero
dx = a / (nx-1)  # Define spatial step
dy = dx
ds = dx
############################
############################
dt = ds / (c * math.sqrt(2))  # Stability Condition
############################
############################

#
## Medium Definition
#
# Coefficients as defined in Taflove,"Compututational Electrodynamics - The Finite
# Difference Time Domain Method", 2nd ed., pg 85.
#
# coefficients subject to the medium parameters % % Field Coefficients

dte = np.ones((nx, ny)) * dt / (ds * eps0)
dtm = np.ones((nx, ny)) * dt / (ds * xmu0)
Da = np.ones((nx, ny))
Db = dtm
Ca = np.ones((nx, ny))
Cb = dte

## Yee's FDTD Algorithm
#
ims = []
for n in range(nt):
    for i in range(nx):
        for j in range(ny):
            if (i == nx2-1 and j == ny2-1):  # Place source in the center
                Ez[i, j] = math.exp(-((n+1 - xn0) / (xndec)) ** 2)  # Gaussian Source
            elif (i == 0 or j == 0 or i == nx-1 or j == ny-1):
                Ez[i, j] = 0  # PEC boundaries at the edges
            else:
                Ez[i, j] = Ez[i, j] * Ca[i, j] + Cb[i, j] * (Hy[i, j] - Hy[i - 1, j]- (Hx[i, j] - Hx[i, j - 1]))

    for i in range(nx):
        for j in range(ny-1):
            Hx[i, j] = Hx[i, j] * Da[i, j] - Db[i, j] * (Ez[i, j + 1] - Ez[i, j])

    for i in range(nx-1):
        for j in range(ny):
            Hy[i, j] = Hy[i, j] * Da[i, j] + Db[i, j] * (Ez[i + 1, j] - Ez[i, j])

    if (n == 20 or n == 60 or n == 80 or n == 100 or n == 120 or n == 140 or n == 160):
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        xd = np.linspace(0, a, 100)
        [xdg, ydg] = np.meshgrid(xd, xd)
        dem3d = ax.plot_surface(xdg, ydg, Ez, cmap='rainbow', edgecolor='none')
        ims.append([dem3d])
        plt.show()

